import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torchvision.transforms as transforms
import argparse
from models.Resnet import ResNet18
from models.Fcanet import make_Fcanet
import matplotlib.pyplot as plt
def train(epoch, model, train_loader, optimizer, criterion, device):
    model.train()
    running_loss = 0.0
    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs, targets = inputs.to(device), targets.to(device)
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()
        if batch_idx % 100 == 99:
            print(f'Epoch [{epoch + 1}], Step [{batch_idx + 1}], Loss: {running_loss / 100:.4f}')
            running_loss = 0.0

def test(model, test_loader, criterion, device):
    model.eval()
    correct = 0
    total = 0
    test_loss = 0.0
    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    accuracy = 100.0 * correct / total
    test_loss /= len(test_loader)
    print(f'Test Accuracy: {accuracy:.2f}%')
    return accuracy, test_loss

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train ResNet or Fcanet on CIFAR-100')
    parser.add_argument('--model', type=int, default=1, help='1 for ResNet, 2 for Fcanet, default is 2')
    parser.add_argument('--strategy', type=str, choices=['LF', 'TS', 'NAS'], default='LF',
                        help='Frequency component selection strategy: LF, TS, or NAS')
    args = parser.parse_args()

    batch_size = 32

    transform = transforms.Compose([
        transforms.RandomCrop(32, padding=4),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize((0.5071, 0.4867, 0.4408), (0.2675, 0.2565, 0.2761)),
    ])

    # 加载数据
    train_dataset = torchvision.datasets.CIFAR100(root='../data', train=True, download=True, transform=transform)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2)

    test_dataset = torchvision.datasets.CIFAR100(root='../data', train=False, download=True, transform=transform)
    test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2)

    # 选择模型
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if device == 'cuda':
        print("using cuda")
    else:
        print("using cpu")
    if args.model == 1:
        model = ResNet18().to(device)
        print("Using the basic ResNet (simplified)")
    elif args.model == 2:
        model = make_Fcanet(args.strategy).to(device)
        print(f"Using the Fcanet with {args.strategy} strategy")
    else:
        model = make_Fcanet('LF').to(device)
        print("Unspecified model! Using the Fcanet with LF strategy")


    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

    best_accuracy = 0.0
    best_model_path = './results/best_model-basic.pth'
    accuracies = []
    for epoch in range(100):
        train(epoch, model, train_loader, optimizer, criterion, device)
        accuracy, test_loss = test(model, test_loader, criterion, device)
        scheduler.step()

        # Save the model if it has the best accuracy so far
        if accuracy > best_accuracy:
            best_accuracy = accuracy
            torch.save(model.state_dict(), best_model_path)
            print(f'Saved the best model with accuracy: {best_accuracy:.2f}%')
        accuracies.append(accuracy)

    plt.figure(figsize=(10, 5))
    plt.plot(range(1, 101), accuracies, label='Test Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title('Test Accuracy over Epochs')
    plt.legend()
    plt.grid(True)
    plt.savefig('accuracy_ResNet.png')
    plt.show()